import pyotp
import qrcode
import time
import json
import os

class TOTPAuthenticator:
    def __init__(self, config_file="accounts.json"):
        self.config_file = config_file
        self.accounts = self._load_accounts()

    def _load_accounts(self):
        """加载已存储的账户密钥[8](@ref)"""
        if os.path.exists(self.config_file):
            with open(self.config_file, 'r') as f:
                return json.load(f)
        return {}

    def _save_accounts(self):
        """保存账户信息到文件[3](@ref)"""
        with open(self.config_file, 'w') as f:
            json.dump(self.accounts, f, indent=4)

    def generate_secret(self, account_name):
        """生成新的TOTP密钥[6](@ref)"""
        secret = pyotp.random_base32()
        self.accounts[account_name] = secret
        self._save_accounts()
        return secret

    def generate_qrcode(self, account_name, issuer="MyAuthApp"):
        """生成绑定二维码[6](@ref)"""
        secret = self.accounts.get(account_name)
        if not secret:
            raise ValueError("Account not found")
        
        uri = pyotp.totp.TOTP(secret).provisioning_uri(
            name=account_name,
            issuer_name=issuer
        )
        img = qrcode.make(uri)
        img.save(f"{account_name}_qrcode.png")
        print(f"二维码已保存为：{account_name}_qrcode.png")

    def get_current_code(self, account_name):
        """获取当前验证码及剩余时间[1](@ref)"""
        secret = self.accounts.get(account_name)
        if not secret:
            return None, 0
        
        totp = pyotp.TOTP(secret)
        remaining = 30 - (int(time.time()) % 30)
        return totp.now(), remaining

    def verify_code(self, account_name, code):
        """验证用户输入的验证码[8](@ref)"""
        secret = self.accounts.get(account_name)
        if not secret:
            return False
        return pyotp.TOTP(secret).verify(code)

def main_menu():
    print("\n=== TOTP身份验证器 ===")
    print("1. 显示所有账户验证码")
    print("2. 添加新账户")
    print("3. 验证验证码")
    print("4. 退出")
    return input("请选择操作：")

# 初始化认证器
auth = TOTPAuthenticator()

while True:
    choice = main_menu()
    
    if choice == '1':
        # 显示所有账户动态码
        for account in auth.accounts:
            code, remaining = auth.get_current_code(account)
            print(f"{account}: {code} (剩余时间: {remaining}秒)")
    
    elif choice == '2':
        # 添加新账户
        account = input("输入账户名称：")
        secret = auth.generate_secret(account)
        auth.generate_qrcode(account)
        print(f"账户已添加！密钥：{secret}")
    
    elif choice == '3':
        # 验证功能
        account = input("账户名称：")
        code = input("输入验证码：")
        if auth.verify_code(account, code):
            print("✅ 验证成功！")
        else:
            print("❌ 验证失败！")
    
    elif choice == '4':
        break

    time.sleep(1)  # 每秒刷新显示
